#import "RegionManager.h"
#import <stdlib.h>
#import <string.h>
#import <mach/mach.h>
#import <mach-o/loader.h>
#import <mach-o/nlist.h>
#import <mach/vm_param.h>

kern_return_t static inline vmRegion(task_t task, 
				     vm_address_t address,
				     Region *region)
{
    vm_offset_t	offset;
    vm_prot_t maxProtection;
    vm_inherit_t inheritance;
    boolean_t shared;
    port_t objectName;
    region->reloc.address = address;
    return vm_region(task,
    		     &region->reloc.address,
		     &region->reloc.size,
		     &region->protection,
		     &maxProtection,
		     &inheritance,
		     &shared,
		     &objectName,
		     &offset);
}

@implementation RegionManager

-initTask: (vm_task_t)theTask readInRegions: (BOOL)readInRegions
{
    vm_statistics_data_t vm_stats;
    
    task = theTask;
    vm_statistics(theTask, &vm_stats);
    pageSize = vm_stats.pagesize;
    relocSize = sizeof(Region);
    rmFlags.shouldSortRelocs = NO;
    if (readInRegions)
	[self readInAllRelocs];
    else
	rmFlags.invalid = YES;
    return self;
}

-(BOOL)isTask
{
    struct task_basic_info taskInfo;
    unsigned int taskInfoCount;
    BOOL ret;
    
    taskInfoCount = TASK_BASIC_INFO_COUNT;
    ret = (task_info(task,
                       TASK_BASIC_INFO,
		       (task_info_t)&taskInfo,
		       &taskInfoCount) == KERN_SUCCESS) ? YES : NO;
    if (!ret && taskGoneCallBack)
	(*taskGoneCallBack)();
    return ret;
}

-(void)setTaskGoneCallBack: (TaskGoneCallBack)theCallBack
{
    taskGoneCallBack = theCallBack;
}

+newTask: (vm_task_t)theTask
{
    return [[super new] initTask: theTask readInRegions: YES];
}

+newTask: (vm_task_t)theTask readInRegions: (BOOL)readInRegions
{
    return [[super new] initTask: theTask readInRegions: readInRegions];
}

-invalidate
{
    if (rmFlags.invalid)
	return self;
    else {
	Region *region;
	int count;
	if (relocs) {
	    for (count = numRelocs, region = (Region *)relocs;
		 count; 
		 count--, region++) {
		if (region->reloc.rFlags.readIn)
		    vm_deallocate(task_self(),
				  region->reloc.data,
				  region->reloc.size);
	    }
	}
	if (pages) {
	    free(pages);
	    pages = NULL;
	}
	return [super invalidate];
    }
}

-(void)getStartPage: (void **)startPage andSize: (int *)sizePage
         forPointer: (void *)start      andSize: (int)numBytes
{
    vm_address_t offset = ((vm_address_t)start % pageSize);
    vm_address_t s = numBytes + offset;
    
    *startPage = start - offset;
    *sizePage =  (s - (s % pageSize)) + pageSize;
}
		
-(void)protectDataAt: (const void *)start for: (int)numBytes
{
    vm_address_t offset = ((vm_address_t)start % pageSize);
    vm_address_t startPage = (vm_address_t)start - offset;
    size_t sizePage = numBytes + offset;
    
    vm_protect(task, startPage, pageSize, NO, VM_PROT_READ);
}

-(void)unProtectDataAt: (const void *)start for: (int)numBytes
{
    vm_address_t offset = ((vm_address_t)start % pageSize);
    vm_address_t startPage = (vm_address_t)start - offset;
    size_t sizePage = numBytes + offset;
    
    vm_protect(task, startPage, pageSize, NO, VM_PROT_READ | VM_PROT_WRITE);
}

-(int)writeDataAt: (const void *)start for: (int)numBytes reloc: (Region *)region markOnly: (BOOL)markOnly
{
    vm_address_t offset = ((vm_address_t)start % pageSize);
    vm_address_t startPage = (vm_address_t)start - offset;
    size_t sizePage = numBytes + offset;
    pointer_t startData;
    kern_return_t ret = 0;
    unsigned char *pPage;
    
    if (markOnly) {
	if (!region->pagesInvalid)
	    regionsInvalid++;
    }
    for (sizePage += (pageSize - (sizePage % pageSize)),
         startData = region->reloc.data
                     + ((pointer_t)startPage - region->reloc.address),
         pPage = region->pages
	         + (((pointer_t)startPage - region->reloc.address) / pageSize);
         sizePage && !ret;
	 sizePage -= pageSize, startData += pageSize, startPage += pageSize,
	 pPage++) {
	if (markOnly) {
	    if (!(*pPage & PAGEINVALID)) {
		*pPage |= PAGEINVALID;
		region->pagesInvalid++;
	    }
	} else {
	    *pPage &= ~PAGEINVALID;
	    if (*pPage & VM_PROT_WRITE)
		ret = vm_write(task, startPage, startData, pageSize);
	    else {
		vm_protect(task, startPage, pageSize, NO, VM_PROT_WRITE);
		ret = vm_write(task, startPage, startData, pageSize);
		vm_protect(task, startPage, pageSize, NO,
			    (vm_prot_t)(*pPage & 0x7));
	    }
	}
    }
    if (ret == KERN_SUCCESS)
	return numBytes;
    else {
	[self isTask];
	return 0;
    }
}

-(int)writeDataAt: (const void *)start for: (int)numBytes reloc: (Region *)region
{
    return [self writeDataAt: start for: numBytes reloc: region markOnly: NO];
}

-(int)putDataAt: (void *)start for: (int)numBytes from: (const void *)data markOnly: (BOOL)markOnly
{
    int numBytesInRegion;
    Reloc *reloc = [self relocFor: start];
    if (reloc) {
        numBytesInRegion = reloc->maxAddress - (int)start;
        if (numBytes > numBytesInRegion)
	    numBytes = numBytesInRegion;
	if (numBytes == 1)
	    *(char *)(reloc->data + ((pointer_t)start - reloc->address))
	    = *(char *)data;
	else
	    memcpy((void *)(reloc->data + ((pointer_t)start - reloc->address)),
	           data,
		   numBytes);
	return [self writeDataAt: start for: numBytes  reloc: (Region *)reloc markOnly: markOnly];
    } else
	return 0;
}

-(int)putDataAt: (void *)start for: (int)numBytes from: (const void *)data
{
    return [self putDataAt: start for: numBytes from: data markOnly: NO];
}

-(void)flushMarkedPages
{
    int ret = 0;
    while (regionsInvalid) {
	int i;
	Region *region;
	
	for (i = 0, region = (Region *)relocs;
	     i < numRelocs;
	     i++, ((void *)region) += relocSize) {
	    if (region->pagesInvalid) {
		unsigned char *p = region->pages;
		while (region->pagesInvalid) {
		    if (*p & PAGEINVALID) {
			int numPage = p - region->pages;
			pointer_t startPage, startData;
			
			startPage
			= region->reloc.address + (numPage * pageSize);
			startData
			= region->reloc.data + (numPage * pageSize);
			*p &= ~PAGEINVALID;
			region->pagesInvalid--;
			if (*p & VM_PROT_WRITE)
			    ret = vm_write(task, startPage, startData, pageSize);
			else {
			    vm_protect(task, startPage, pageSize, NO, VM_PROT_WRITE);
			    ret = vm_write(task, startPage, startData, pageSize);
			    vm_protect(task, startPage, pageSize, NO,
					(vm_prot_t)(*p & 0x7));
			}
		    }
		    p++;
		}
		regionsInvalid--;
	    }
	}
    }
    if (ret) 
	[self isTask];
}

-(int)writeDataAt: (const void *)start for: (int)numBytes;
{
    Reloc *reloc = [self relocFor: start];
    if (reloc)
	return [self writeDataAt: start for: numBytes reloc: (Region *)reloc];
    else
	return NO;
}

-(BOOL)readInReloc: (Reloc *)reloc
{
    kern_return_t ret;
    unsigned int nData;
    if (reloc->rFlags.readIn) {
	return YES;
    } else {
	ret = vm_read(task, reloc->address, reloc->size, &reloc->data, &nData);
	if (ret == KERN_SUCCESS) {
	    reloc->maxData = reloc->data + reloc->size;
	    reloc->displacement = reloc->data - reloc->address;
	    reloc->rFlags.readIn = YES;
	    return YES;
	} else {
	    [self isTask];
	    return NO;
	}
    }
}

-(void)combineRegions: (BOOL)combineRegions
{ dontCombineRegions = !combineRegions;}

-(void)_readInAllRelocs
{
    kern_return_t error = 0;
    vm_address_t address = VM_MIN_ADDRESS;
    Region *theRegion, *regions, *newRegion, *nextRegion, *lastRegion;
    int count = 0, numPages = 0, numAlloced = 16, nPages;
    unsigned char *thePage;
    
    if ([self isTask]) {
	theRegion = regions = malloc(numAlloced * sizeof(Region));
	bzero(regions, numAlloced * sizeof(Region));
	while (!error) {
	    error = vmRegion(task, address, theRegion);
	    address = theRegion->reloc.address + theRegion->reloc.size;
	    if (!error) {
		if (theRegion->protection & VM_PROT_READ) {
		    count++;
		    numPages += theRegion->reloc.size / pageSize;
		    if (count < numAlloced)
			theRegion++;
		    else {
			numAlloced *= 2;
			regions = realloc(regions,
			                  numAlloced * sizeof(Region));
			theRegion = regions + count;
			bzero(theRegion, count * sizeof(Region));
		    }
		}
		if (address < theRegion->reloc.address)
		    error = KERN_NO_SPACE;
	    }
	}
	if (error == KERN_NO_SPACE) {
	    lastRegion = theRegion;
	
	    pages = malloc(numPages);
	    for (theRegion = regions, thePage = pages;
	         theRegion < lastRegion;
		 theRegion++) {
		theRegion->pages = thePage;
		nPages = theRegion->reloc.size / pageSize;
		while(nPages--)
		    *(thePage++) = (unsigned char)theRegion->protection;
	    }
	    
	    if (dontCombineRegions) {
		for (theRegion = regions; theRegion < lastRegion; theRegion++)
		    theRegion->reloc.maxAddress
		    = theRegion->reloc.address + theRegion->reloc.size;
		relocs = regions;
		numRelocs = count;
	    } else {
		relocs = malloc(count * sizeof(Region));
		bzero(relocs, count * sizeof(Region));
		for (theRegion = regions, newRegion = relocs;
		    theRegion < lastRegion;
		    theRegion = nextRegion, newRegion++) {
		    for (nextRegion = theRegion + 1;
			(nextRegion < lastRegion
			&& ((theRegion->reloc.address + theRegion->reloc.size)
			    == nextRegion->reloc.address));
			nextRegion++) {
			theRegion->reloc.size += nextRegion->reloc.size;
		    }
		    *newRegion = *theRegion;
		    newRegion->reloc.maxAddress
		    = newRegion->reloc.address + newRegion->reloc.size;
		}
		numRelocs = newRegion - (Region *)relocs;
		relocs = realloc(relocs, numRelocs * sizeof(Region));
		free(regions);
	    }
	    rmFlags.invalid = NO;
	} else
	    free(regions);
	    
	if (error != KERN_NO_SPACE) 
	    [self isTask];
    } 
}

-(struct mach_header *)getMachHeader
{
    Reloc *reloc;
    int count;
    struct mach_header *header, *foundHeader;

    if (rmFlags.invalid)
	[self readInAllRelocs];
    reloc = relocs;
    for (foundHeader = NULL, count = 0;
         !foundHeader && count < numRelocs;
	 count++) {
	[self readInReloc: reloc];
	header = [self pointerFor: (void *)reloc->address
	                 withSize: sizeof(*header)];
	if (header
	    && (header->magic == MH_MAGIC)
	    && ([self pointerFor: (void *)reloc->address
	                withSize: header->sizeofcmds]))
	    foundHeader = header;
	else
	    ((void *)reloc) += relocSize;
    }
    return foundHeader;
}

-(int)getNumMachHeaders
{
    struct mach_header *myHeader;
    int numHeaders, i;
    struct fvmlib_command *loadCmd;
    if (myHeader = [self getMachHeader]) {
	numHeaders = 1;
	for (i = 0, loadCmd = (struct fvmlib_command *)(myHeader + 1);
	     i < myHeader->ncmds;
	     i++,
	     loadCmd = ((void *)loadCmd + loadCmd->cmdsize)) {
	    if (loadCmd->cmd == LC_LOADFVMLIB)
		numHeaders++;
	}
    } else
	numHeaders = 0;
    return numHeaders;
}

-(struct mach_header **)getMachHeadersWithNames: (char ***)names
{
    struct mach_header *myHeader, **headers, *header;
    char **theNames = NULL;
    struct fvmlib_command *loadCmd;
    int numHeaders, headerIndex;
    numHeaders = [self getNumMachHeaders];
    myHeader = [self getMachHeader];
    headers = malloc((numHeaders + 1) * sizeof(*headers));
    if (names) {
	*names = theNames = malloc((numHeaders + 1) * sizeof(*theNames));
	theNames[0] = NULL;
    }
    headers[0] = myHeader;
    for (headerIndex = 1, loadCmd = (struct fvmlib_command *)(myHeader + 1); 
	 headerIndex < numHeaders;
	 loadCmd = ((void *)loadCmd + loadCmd->cmdsize)) {
	if (loadCmd->cmd == LC_LOADFVMLIB) {
	    header = [self pointerFor: (void *)(loadCmd->fvmlib.header_addr) 
	                     withSize: sizeof(*header)];
	    if (header)
		headers[headerIndex]
		= [self pointerFor: (void *)(loadCmd->fvmlib.header_addr)
	                  withSize: header->sizeofcmds];
	    if (names)
		theNames[headerIndex]
		= (char *)loadCmd
			    + ((struct fvmlib_command *)loadCmd)
			    ->fvmlib.name.offset; 
	    headerIndex++;
	}
    }
    headers[numHeaders] = NULL;
    if (names)
	theNames[numHeaders] = NULL;
    return headers;
}

-(struct mach_header **)getMachHeaders
{
    return [self getMachHeadersWithNames: NULL];
}

-(unsigned)getSlide
{
	return 0;
}

-(struct mach_header*)copyLoadedImage:(struct mach_header*)header_addr
		      withSlide:(unsigned)slide
{
  struct load_command *cmd;
  unsigned long fileSize = 0;
  struct mach_header *header, *newHeader;
  struct symtab_command *symtabCmd = 0;
  kern_return_t r;
  int i;
  unsigned header_size;


  header = [self pointerFor:header_addr 
		 withSize:sizeof (struct mach_header)];
  header_size = header->sizeofcmds + sizeof (struct mach_header);
  
  header = [self pointerFor:header_addr
		 withSize:header_size];

  // get the size of the original file
  for (cmd = (struct load_command *)(header + 1), i = 0;
       i < header->ncmds;
       ((void *)cmd) += cmd->cmdsize, i++) 
    {
      unsigned long topOff;

      if (cmd->cmd == LC_SEGMENT)
	{
	  struct segment_command *segCmd = (struct segment_command *)cmd;
	  topOff = segCmd->fileoff + segCmd->filesize;
	  
	  if (topOff > fileSize)
	    fileSize = topOff;
	}
      else if (cmd->cmd == LC_SYMTAB)
	{
	  symtabCmd = (struct symtab_command*)cmd;
	}
    }

  r = vm_allocate (task_self (), (vm_address_t*)&newHeader, fileSize, 1);
  if (r != KERN_SUCCESS) [self error:"vm_allocate failed %d", r];

  // get the size of the original file
  for (cmd = (struct load_command *)(header + 1), i = 0;
       i < header->ncmds;
       ((void *)cmd) += cmd->cmdsize, i++)
    {
      if (cmd->cmd == LC_SEGMENT)
        {
          struct segment_command *segCmd = (struct segment_command *)cmd;
          void *segment = [self pointerFor:(char*)segCmd->vmaddr + slide
                                  withSize:segCmd->filesize];

          // get size rounded up to page size...
          vm_size_t size = segCmd->filesize;
          if (size % vm_page_size != 0)
            {
              size += vm_page_size - (size % vm_page_size);
            }

          if ((segCmd->fileoff % vm_page_size) == 0)
            {
              vm_copy (task_self (),
                       (vm_address_t)segment,
                       (vm_size_t)size,
                       (vm_address_t)((void*)newHeader + segCmd->fileoff));
            }
          else
            {
              memcpy (segment, (void*)newHeader + segCmd->fileoff, size);
            }
        }
    }

  // slide the image by modifying the load commands
  for (cmd = (struct load_command *)(newHeader + 1), i = 0;
       i < header->ncmds;
       ((void *)cmd) += cmd->cmdsize, i++) 
    {
      if (cmd->cmd == LC_SEGMENT)
      {
	struct segment_command *segCmd = (struct segment_command *)cmd;

	{
	  int n;
	  struct section *sect;
	  segCmd->vmaddr += slide;
	  
	  for (sect = (struct section*) (segCmd + 1), n = 0; 
	       n < segCmd->nsects; 
	       sect += 1, n += 1)
	    {
	      sect->addr += slide;
	    }
	}
      }
    }

  // now relocate the symbol table information
  {
    struct nlist* symtab = (struct nlist*) ((char*)newHeader 
					    + symtabCmd->symoff);
    for (i = 0; i < symtabCmd->nsyms; i++)
      {
	if (symtab[i].n_sect != NO_SECT)
	  {
	    symtab[i].n_value += slide;
	  }
      }
  }

  return newHeader;
}

@end
